"""
    ================================================================================
                            ------------utf-8--------------
    ================================================================================
@Author: 
    rfdsg
@Create Time: 
    2025/5/20 - 18:29
@Description:

@Attention:
    
"""
import inspect
import logging
from typing import Callable, Optional, Union

import networkx as nx

from factor_lib.common_tool import DS, DSN


class Graph:
    registry_funcs = {}

    def __init__(self, external_inputs: dict[str, DSN]):
        """
        external_inputs: dict，外部输入数据，如 {'factor': xxx, 'returns': yyy}
        """
        self.graph = nx.DiGraph()
        self.values: dict[str, dict[str, Optional[Union[DSN, Callable]]]] = {}

        # 先把外部输入加入values，标记为“已计算”
        if external_inputs:
            for k, v in external_inputs.items():
                self.values[k] = {"func": None, "result": v}

        # 注册函数，自动从参数获取依赖（跳过self）
        for name, func in self.registry_funcs.items():
            sig = inspect.signature(func)
            depends = [p.name for p in sig.parameters.values() if p.name != "self"]
            self._register(name, func, depends)

    @classmethod
    def evaluator(cls, func):
        cls.registry_funcs[func.__name__] = func
        return func

    def _register(self, name: str, func: Callable, depends=()):
        self.graph.add_node(name)
        for dep in depends:
            self.graph.add_edge(dep, name)
        self.values[name] = {"func": func, "result": None}

    def register_const(self, name: str, value):
        """动态注册外部输入，并清空依赖这个输入的缓存结果"""
        self.values[name] = {"func": None, "result": value}
        self.graph.add_node(name)  # 防止是新节点

    def evaluate(self, name: str, debug: bool = False):
        if name not in self.values:
            raise KeyError(f"没有这样一个评估函数叫{name}")

        if self.values[name]["result"] is not None:
            return self.values[name]["result"]
        # 如果是外部输入，没有函数，直接返回结果
        if self.values[name]["func"] is None:
            return self.values[name]["result"]
        # 先递归计算依赖
        dep_list = [dep for dep in self.graph.predecessors(name)]
        if debug:
            logging.info(f"[args获取] {name} <- {dep_list}")
        args = [self.evaluate(dep, debug=debug) for dep in dep_list]
        if debug:
            logging.info(f"[开始计算] {name} <- {dep_list}")
        # 注意传入self，执行实例方法
        result = self.values[name]["func"](self, *args)
        self.values[name]["result"] = result
        return result

    def clear(self):
        for val in self.values.values():
            val["result"] = None

    def all_results(self):
        return {k: v["result"] for k, v in self.values.items()}

    def evaluate_all(self, debug: bool = False):
        results = {}
        for name, info in self.values.items():
            # 只处理有函数的节点（即评估函数），不包括外部输入
            if info["func"] is not None:
                results[name] = self.evaluate(name, debug=debug)
        return results
